package com.ams.io;

import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.ConcurrentLinkedQueue;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class BufferAllocator implements IBufferAllocator {
    final private Logger logger = LoggerFactory
            .getLogger(BufferAllocator.class);

    private ByteBuffer byteBuffer = null;

    private int chunkSize = 2 * 1024 * 1024; // 2M

    private int poolSize = 64 * 1024 * 1024; // 64M

    private ConcurrentLinkedQueue<ByteBuffer> chunkPool = new ConcurrentLinkedQueue<ByteBuffer>();

    public void init() {
        byteBuffer = null;
        for (int i = 0; i < poolSize / chunkSize; i++) {
            ByteBuffer buf = allocateBuffer(chunkSize);
            chunkPool.offer(buf);
        }
        ByteBufferCollector collector = new ByteBufferCollector();
        collector.start();
    }

    private ByteBuffer allocateBuffer(int size) {
        return ByteBuffer.allocateDirect(size);
    }

    private ByteBuffer newBuffer(int size) {
        ByteBuffer chunk = null;
        if (size > chunkSize) {
            chunk = allocateBuffer(size);
            logger.debug("allocate direct buffer:{}", size);
        } else {
            chunk = chunkPool.poll();
            if (chunk == null) {
                chunk = allocateBuffer(chunkSize);
                logger.debug("allocate direct buffer:{}", chunkSize);
            }
        }
        ByteBuffer chunkHolder = chunk.duplicate();
        referenceList.add(new ChunkReference(chunkHolder, chunk));
        return chunkHolder;
    }

    private class ChunkReference extends WeakReference<ByteBuffer> {
        private ByteBuffer chunk;

        public ChunkReference(ByteBuffer referent, ByteBuffer chunk) {
            super(referent, chunkReferenceQueue);
            this.chunk = chunk;
        }

        public ByteBuffer getChunk() {
            return chunk;
        }
    };

    private static ReferenceQueue<ByteBuffer> chunkReferenceQueue = new ReferenceQueue<ByteBuffer>();
    private static List<ChunkReference> referenceList = Collections
            .synchronizedList(new LinkedList<ChunkReference>());

    private class ByteBufferCollector extends Thread {
        private static final int COLLECT_INTERVAL_MS = 1000;

        private void recycle(ByteBuffer buf) {
            buf.clear();
            chunkPool.offer(buf);
        }

        private void collect() {
            ChunkReference ref;
            int n = 0;
            while ((ref = (ChunkReference) chunkReferenceQueue.poll()) != null) {
                ByteBuffer chunk = ref.getChunk();
                recycle(chunk);
                referenceList.remove(ref);
                n++;
            }
            if (n > 0) {
                logger.debug("collected {} chunk buffers", n);
            }

        }

        public ByteBufferCollector() {
            super();
            try {
                setDaemon(true);
            } catch (Exception e) {
            }
        }

        public void run() {
            try {
                while (!Thread.interrupted()) {
                    sleep(COLLECT_INTERVAL_MS);
                    collect();
                }
            } catch (InterruptedException e) {
                interrupt();
            }
        }
    }

    public synchronized ByteBuffer allocate(int size) {
        if (byteBuffer == null
                || byteBuffer.capacity() - byteBuffer.limit() < size) {
            byteBuffer = newBuffer(size);
        }
        byteBuffer.limit(byteBuffer.position() + size);
        ByteBuffer slice = byteBuffer.slice();
        byteBuffer.position(byteBuffer.limit());

        return slice;
    }

    public void setPoolSize(int poolSize) {
        this.poolSize = poolSize;
    }

    public void setChunkSize(int chunkSize) {
        this.chunkSize = chunkSize;
    }

}
